import time
t1=time.perf_counter()
from langgraph.graph import StateGraph, END
from typing import TypedDict
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(
        model="deepseek-chat",
        base_url="https://api.deepseek.com/v1",
        temperature=0.7,
        api_key="key"
)

start_time=time.perf_counter()
class AgentState(TypedDict):
    plan_result: str
    topic:str
    chapter_content:str
    story_content:str
    scores:str
plan_begin=time.perf_counter()
def plan(state:AgentState):
    topic=state['topic']
    prompt=f'''
    ##角色
    你是一位文章创意的专家,构思以{topic}为要求的文章内容创意。
    按照{topic}，你需要准备文章的一个详细的大纲，按照文章类型可以适当增添内容到大纲中。
      你需要构思12个主要章节，章节内容要详细，第13章为感悟，章节小标题以一种吸引人的方式命名。
      你的任务是根据用户提供的要求，提供适合读者阅读的、专业的、有见地的文章内容创意。
      1.文章结构清晰，富有想象力，具有创新性，主题明确有深度，吸引读者阅读。
      2.文章语言表达与风格同样重要，文笔流畅、具有独特的语言风格以及恰当的修辞。
      3.文章有一定的教育意义，遵守道德和法律准则。
      4.构思的文章章节应为12个和一个感悟总结章节。
      5.文章鲜明，恰当的修辞和风格、张力与节奏以及有意外与惊喜。
      6.文章能引起读者情感共鸣，注意其在创意、语言描述等方面的协调性。
      大纲以Markdown形式输出。
      
    ##用户
    用户要求：{topic}'''
    prompt=PromptTemplate.from_template(prompt)
    prompt=prompt.format(topic=topic)
    plan_result=llm.invoke(prompt).content
    state['plan_result']=plan_result
    return {'plan_result':plan_result}

plan_final=time.perf_counter()
plan_time=plan_final-plan_begin

write_begin=time.perf_counter()
chapter_number = 1
story_content = ""
def write(state:AgentState):
    plan_result=state['plan_result']
    global chapter_number
    global story_content
    prompt=f'''
     ##角色
     你是一位文章章节编写专家，你编写的章节内容有创新性，特别适合用户阅读。
     1.根据{plan_result}的大纲内容，可以按照需求适当添加细节，扩充不少于800个汉字的第{chapter_number}个章节内容。
      2.第{chapter_number}个章节段落内容不能重复，应有恰当的修辞以让字数不少于800个汉字。
      3.第二章根据{story_content}的第一章章节发展和{plan_result}的第二章续写，以此类推，直至第十二章。
      4.第一至十二章不用写感悟。
      5.根据{plan_result}扩充第十三章且字数不少于500字。
      6.每一章节要有标题和标注第几章，第一章最前面先输出文章标题，以Markdown形式输出
      7.只需要输出第{chapter_number}章节内容，不需要输出其他章节。
      开头输出文章标题并且居中显示，章节要有标题和标注第几章，以Markdown形式输出。
      
      ##用户
    大纲：{plan_result}
    章节：{chapter_number}'''
    prompt = PromptTemplate.from_template(prompt)
    prompt=prompt.format(plan_result=plan_result,chapter_number=chapter_number)
    chapter_content=llm.invoke(prompt).content
    chapter_number +=1
    story_content +=chapter_content + "\n"
    state['story_content'] = story_content
    return {'story_content':state['story_content']}

write_final=time.perf_counter()
write_time=write_final-write_begin

score_begin=time.perf_counter()
def score(state:AgentState):
    story_content=state['story_content']
    prompt=f'''
     ##角色
     你是一位文章打分专家，为了让文章更符合要求，你需要对文章内容{story_content}进行打分。
      你的任务是根据
      1.评估文章的创意(5分)与主题(5分)，考察其创新性(5分)、主题深度(5分)以及明确性(5分)。
      2.深入分析结构构建(10分)，包括逻辑性(10分)、张力与节奏(5分)以及意外与惊喜(5分)。
      3.我们关注文章内容的丰富性(5分)，无明显的错误(10分)以及语言的准确性(5分)。
      4.语言表达与风格同样重要，文笔的流畅性(5分)、独特的语言风格(5分)以及恰当的修辞(5分)。
      5.综合考虑文章的整体印象，评估其在创意、结构、文章风格等方面的协调性(5分)，以及是否能够触动读者的心灵，留下深刻印象(5分)
      6.计算这次评分最后的总分。
      7.按照标准进行10次评分，计算出的10次评分总分相加的平均值为最后的分数scores例如： 3并scores在最后一行以数字形式输出。
      标准进行打分，计算出的10次评分的平均值即为最终的分数，例如：3。
      
      ##用户
    故事：{story_content}'''
    prompt=PromptTemplate.from_template(prompt)
    prompt=prompt.format(story_content=story_content)
    scores=llm.invoke(str(prompt)).content
    state['scores']=str(scores)
    return {'scores':scores}

score_final=time.perf_counter()
score_time=score_final-score_begin

time1=time.perf_counter()
def rewrite(state: AgentState):
    if chapter_number<13:
        return "rewrite"
    else:
        return "no_write"
    
workflow = StateGraph(AgentState)
# Define the nodes and their execution order
workflow.add_node("plan", plan)
workflow.add_node("write", write)
workflow.add_node("score", score)
# Set up the flow
workflow.set_entry_point("plan")
workflow.add_edge("plan", "write")
workflow.add_conditional_edges(
    "write",
    rewrite,
    {
        "rewrite": "write",
        "no_write": "score",
    },
)
workflow.add_edge("score", END)
# Compile and run the workflow
app = workflow.compile()
topic=input("用户要求")
r1 = str(app.invoke({'topic': topic}).get('plan_result'))
r2 = str(app.invoke({'topic': topic}).get('story_content'))
r3 = str(app.invoke({'topic': topic}).get('scores'))
time2=time.perf_counter()
result = r2 +'\n'+ r3
finally_time=time.perf_counter()

cooperation=round(((plan_time+write_time+score_time+time2-time1)/(finally_time-start_time))*100,2)
cost=round((len(r1)+len(r2)+len(r3))*2*0.000002,2)
score_weight = 0.5
resource_weight = 0.1
collaboration_weight = 0.4
import re
sco=r3[-9:]
match = re.search(r'\d+', sco)
if match:
    score_value = int(match.group())
else:
    print("无法从 score_result 中提取数字")
success_rate=round((score_value * score_weight + cost * resource_weight + cooperation * collaboration_weight),2)
print(result)
file = open(r"..\lang\file3.txt", "w", encoding="utf-8")
file.write(story_content)
file.close()
file = open(r"..\lang\score3.txt", "w", encoding="utf-8")
file.write(r3)
file.close()
print("耗费时间："+str(finally_time-t1))
print("协作效率："+str(cooperation)+"%")
print("花费："+str(cost))
print("综合评分："+str(success_rate)+"%")
file = open(r"..\lang\result3.txt", "w", encoding="utf-8")
file.write("耗费时间："+str(finally_time-t1) +'\n'+ "协作效率："+str(cooperation)+"%" +'\n'+ "花费："+str(cost) +'\n'+ "综合评分："+str(success_rate)+"%")
file.close()